package com.example.flinkcourse.lesson4.state;

import org.apache.flink.api.common.restartstrategy.RestartStrategies;
import org.apache.flink.api.common.time.Time;
import org.apache.flink.contrib.streaming.state.RocksDBStateBackend;
import org.apache.flink.runtime.state.StateBackend;
import org.apache.flink.streaming.api.CheckpointingMode;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.concurrent.TimeUnit;

/**
 * 状态管理器
 * 用于配置状态后端和检查点
 */
public class StateManager {
    private static final Logger LOG = LoggerFactory.getLogger(StateManager.class);

    /**
     * 配置状态后端
     * @param env 流执行环境
     * @param checkpointDir 检查点目录
     * @param incrementalCheckpoints 是否启用增量检查点
     */
    public static void configureStateBackend(
            StreamExecutionEnvironment env,
            String checkpointDir,
            boolean incrementalCheckpoints) {
        LOG.info("Configuring state backend...");
        
        try {
            // 配置 RocksDB 状态后端
            RocksDBStateBackend stateBackend = new RocksDBStateBackend(
                checkpointDir,
                incrementalCheckpoints
            );
            env.setStateBackend(stateBackend);
            
            LOG.info("State backend configured successfully");
        } catch (Exception e) {
            LOG.error("Error configuring state backend", e);
            throw new RuntimeException("Failed to configure state backend", e);
        }
    }
    
    /**
     * 配置检查点
     * @param env 流执行环境
     * @param interval 检查点间隔（毫秒）
     * @param timeout 检查点超时（毫秒）
     * @param minPauseBetweenCheckpoints 检查点之间的最小暂停时间（毫秒）
     * @param maxConcurrentCheckpoints 最大并发检查点数量
     */
    public static void configureCheckpoint(
            StreamExecutionEnvironment env,
            long interval,
            long timeout,
            long minPauseBetweenCheckpoints,
            int maxConcurrentCheckpoints) {
        LOG.info("Configuring checkpoint...");
        
        // 启用检查点
        env.enableCheckpointing(interval);
        
        // 设置检查点模式
        env.getCheckpointConfig().setCheckpointingMode(CheckpointingMode.EXACTLY_ONCE);
        
        // 设置检查点超时
        env.getCheckpointConfig().setCheckpointTimeout(timeout);
        
        // 设置检查点之间的最小暂停时间
        env.getCheckpointConfig().setMinPauseBetweenCheckpoints(minPauseBetweenCheckpoints);
        
        // 设置最大并发检查点数量
        env.getCheckpointConfig().setMaxConcurrentCheckpoints(maxConcurrentCheckpoints);
        
        LOG.info("Checkpoint configured successfully");
    }
    
    /**
     * 配置重启策略
     * @param env 流执行环境
     * @param maxAttempts 最大重启次数
     * @param delayBetweenAttempts 重启延迟（毫秒）
     */
    public static void configureRestartStrategy(
            StreamExecutionEnvironment env,
            int maxAttempts,
            long delayBetweenAttempts) {
        LOG.info("Configuring restart strategy...");
        
        // 设置固定延迟重启策略
        env.setRestartStrategy(
            RestartStrategies.fixedDelayRestart(
                maxAttempts,
                Time.of(delayBetweenAttempts, TimeUnit.MILLISECONDS)
            )
        );
        
        LOG.info("Restart strategy configured successfully");
    }
    
    /**
     * 配置默认状态后端和检查点
     * @param env 流执行环境
     */
    public static void configureDefault(StreamExecutionEnvironment env) {
        LOG.info("Configuring default state backend and checkpoint...");
        
        // 配置状态后端
        configureStateBackend(env, "hdfs:///flink/state", true);
        
        // 配置检查点
        configureCheckpoint(env, 60000, 30000, 10000, 1);
        
        // 配置重启策略
        configureRestartStrategy(env, 3, 10000);
        
        LOG.info("Default configuration completed successfully");
    }
} 